#include <TApplication.h>
#include <TGClient.h>
#include <TGButton.h>
#include <TGNumberEntry.h>
#include <TGComboBox.h>
#include <TRootEmbeddedCanvas.h>
#include <TGLabel.h>
#include <TCanvas.h>
#include <TPad.h>
#include <TGraph.h>
#include <TLegend.h>
#include <TTree.h>
#include <TFile.h>
#include <TClonesArray.h>
#include <TSystem.h>
#include <vector>
#include <iostream>
#include <fairlogger/Logger.h>

class DataViewer : public TGMainFrame {
private:
    TRootEmbeddedCanvas* fCanvas;
    TGNumberEntry* fTimeSliceIDEntry; //输入框：选择时间片编号
    TGComboBox* fPXICombo; //下拉框：选择PXI编号
    TGLabel* fTimeSliceLabel; //打印时间片的值

    TFile* fFile{nullptr};
    TTree* fTree{nullptr};
    TBranch* fBranch{nullptr};
    TClonesArray* fData{nullptr};

    int fCurrentTimeSliceIndex = 0; // 当前选择的时时间片索引
    int fCurrentPXIIndex = 0; // 当前选择的PXI索引

    struct TimeSliceInfo {
        int entry;  // 此时间片对应 TTree 的 entry
        std::map<
            uint8_t,  // PXIID
            std::map<
                uint8_t,  // FEEID
                std::map<
                    uint8_t,        // ChannelID
                    std::vector<int> // 对应该 FEE/Channel 在本 entry 中的TClonesArray索引
                >
            >
        > pxiMap;
    };
    std::map<uint64_t, TimeSliceInfo> fIndexMap;
    std::map<uint64_t, TimeSliceInfo>::iterator fItTS;

public:
    DataViewer(const TGWindow* p, UInt_t w, UInt_t h, const char* fname);
    virtual ~DataViewer();

    void UpdateTimeSliceCombo(); // 更新PXI下拉框
    void OnPXISelected(Int_t id); // PXI选择事件
    void DoDraw(); // 画图按钮事件


ClassDef(DataViewer, 0);

};

DataViewer::DataViewer(const TGWindow* p, UInt_t w, UInt_t h, const char* fname)
    : TGMainFrame(p, w, h)
{
    SetCleanup(kDeepCleanup);

    // Load file and tree
    fFile = TFile::Open(fname);
    if (!fFile || fFile->IsZombie()) {
        LOG(error) << "Cannot open file.";
        return;
    }
    fTree = (TTree*) fFile->Get("HiTree");
    fTree->SetBranchAddress("HiTpcUnpackedData", &fData);
    Long64_t nEntries = fTree->GetEntries();
    /***********初始化索引map**************/
    for (Long64_t i = 0; i < nEntries; i++) {
        /**********进度条**********/
        int barWidth = 50;   // 进度条宽度
        double progress = double(i) / nEntries;
        std::cout << "\r[";
        int pos = barWidth * progress;
        for (int k = 0; k < barWidth; k++) {
            if (k < pos) std::cout << "=";
            else if (k == pos) std::cout << ">";
            else std::cout << " ";
        }
        std::cout << "] " << int(progress * 100.0) + 1 << "%  ("
                  << i + 1<< "/" << nEntries << ")";
        std::cout.flush();
        /****************************/
        fTree->GetEntry(i);
        Long64_t nTClonesEntries = fData->GetEntriesFast();
        for (int j = 0; j < nTClonesEntries; j++) {
            auto* data = (HiTpcUnpackedData*)fData->At(j);
            if(j == 0){
                fIndexMap[data->GetTimeSliceID()].entry = i;
            }
            fIndexMap[data->GetTimeSliceID()].pxiMap[data->GetPxiID()][data->GetFeeID()]
                [data->GetChannelID()].emplace_back(j);
            LOG(debug)<<"TS = "<<data->GetTimeSliceID()<<", PXI = "<<int(data->GetPxiID())
            <<", FEE = "<<int(data->GetFeeID())<<", ch = "<<int(data->GetChannelID());
        }
    }
    std::cout << "\nDone!\n";
    /***********输入信息控制区*************/
    int Width = 10;
    int left = 5, right = 5, top = 5, bottom = 5;
    int ComboX = 80, ComboY = 20;
    TGHorizontalFrame* ctrlFrame = new TGHorizontalFrame(this, w, 30);
    AddFrame(ctrlFrame, new TGLayoutHints(kLHintsTop | kLHintsExpandX));
    
    //时间片部分
    TGLabel* TimeSliceIDSet = new TGLabel(ctrlFrame, "TimeSlice ID Set: ");
    //四个数字代表相对于左右上下边框的距离
    ctrlFrame->AddFrame(TimeSliceIDSet, 
        new TGLayoutHints(kLHintsTop | kLHintsLeft, left, right, top, bottom));
    fTimeSliceIDEntry = new TGNumberEntry(ctrlFrame, 0, Width);
    fTimeSliceIDEntry->SetLimits(TGNumberFormat::kNELLimitMinMax, 0, fIndexMap.size() - 1);
    ctrlFrame->AddFrame(fTimeSliceIDEntry, 
        new TGLayoutHints(kLHintsTop | kLHintsLeft, left, right, top, bottom));
    
    //PXI部分    
    TGLabel* PxiLabel = new TGLabel(ctrlFrame, "PXI Value Set: ");
    ctrlFrame->AddFrame(PxiLabel, 
        new TGLayoutHints(kLHintsTop | kLHintsLeft, left, right, top, bottom));
    fPXICombo = new TGComboBox(ctrlFrame);
    fPXICombo->Resize(ComboX, ComboY);
    fPXICombo->Connect("Selected(Int_t)", "DataViewer", this, "OnPXISelected(Int_t)");
    ctrlFrame->AddFrame(fPXICombo, 
        new TGLayoutHints(kLHintsTop | kLHintsLeft, left, right, top, bottom));
    
    //时间片部分
    fTimeSliceLabel = new TGLabel(ctrlFrame, "TimeSlice Value: ");
    ctrlFrame->AddFrame(fTimeSliceLabel, 
        new TGLayoutHints(kLHintsTop | kLHintsLeft, left, right, top, bottom));

    /***********************************/


    /***************画图区****************/
    fCanvas = new TRootEmbeddedCanvas("Canvas", this, w, h);
    AddFrame(fCanvas, new TGLayoutHints(kLHintsExpandX | kLHintsExpandY));
    fTimeSliceIDEntry->Connect("ValueSet(Long_t)", "DataViewer", this, "UpdateTimeSliceCombo()");
    /***********************************/

    SetWindowName("Unpacked Data Viewer");
    MapSubwindows();
    Resize();
    MapWindow();
    
    // 读取数据后动态填充下拉框
    UpdateTimeSliceCombo();

}

DataViewer::~DataViewer() {
    Cleanup();
    if (fFile) fFile->Close();
}

void DataViewer::UpdateTimeSliceCombo() {
    fCurrentTimeSliceIndex = fTimeSliceIDEntry->GetNumber();
    fItTS = std::next(fIndexMap.begin(), fCurrentTimeSliceIndex);
    fTree->GetEntry(fItTS->second.entry);

    fTimeSliceLabel->SetText(Form("TimeSlice Value Real: %lu ", fItTS->first));

    fPXICombo->RemoveAll();
    for (auto& [pxi, _] : fItTS->second.pxiMap) {
        fPXICombo->AddEntry(Form("PXI %d", pxi), pxi);
    }

    // --- 如果当前选中的PXI依然存在，就保持选中 ---
    if (fItTS->second.pxiMap.count(fCurrentPXIIndex)){
        fPXICombo->Select(fCurrentPXIIndex, false);
        OnPXISelected(fCurrentPXIIndex);
    } 
    else {  // 否则默认选第一个
        fCurrentPXIIndex = fItTS->second.pxiMap.begin()->first;
        fPXICombo->Select(fCurrentPXIIndex, false);
        OnPXISelected(fCurrentPXIIndex);
    }
}

void DataViewer::OnPXISelected(Int_t id) {
    fCurrentPXIIndex = id; 
    DoDraw();
}

void DataViewer::DoDraw(){

    TCanvas* c = fCanvas->GetCanvas();
    c->Clear();
    c->Divide(2, 2);  // 

    //初始化波形数据
    std::vector<double> x(4096, 0), y(4096, 0);
    for (size_t i = 0; i < 4096; ++i) {
        x[i] = i * 0.025; // 单位 us
    }

    for(auto& [fee, channelMap] : fItTS->second.pxiMap[fCurrentPXIIndex]){
        c->cd(fee + 1);
        gPad->SetGrid();
        TLegend* legend = new TLegend(0.7, 0.7, 0.9, 0.9);
        legend->SetNColumns(3);
        // 使用 TMultiGraph 管理所有 cluster 的显示范围
        TMultiGraph* mg = new TMultiGraph();
        int index = 0;
        for(auto& [channel, entry] : channelMap){
            std::fill(y.begin(), y.end(), 0); // 清空旧波形
            for(auto& iter : entry){
                auto* data = (HiTpcUnpackedData*)fData->At(iter);
                for (size_t j = 0; j < data->GetWaveformData().size(); j++) {
                    y[data->GetTimeOffset() + j] = data->GetWaveformData()[j]; 
                }
            }
            auto* g = new TGraph(x.size(), x.data(), y.data());
            int colorIdx = (int)((1.0 * index / channelMap.size()) * 
                (gStyle->GetNumberOfColors() - 1));
            g->SetLineColor(TColor::GetColorPalette(colorIdx));
            g->SetLineWidth(2);
             
            g->SetMarkerStyle(6);
            mg->Add(g);
            legend->AddEntry(g, Form("Channel %d", channel), "lp");
            index++;
        }
        mg->SetTitle(Form("FEE %d;Time [us];Amplitude", fee));
        mg->GetYaxis()->SetTitleOffset(1.1);
        mg->GetYaxis()->CenterTitle(); 
        mg->Draw("APL");
        legend->SetTextSize(0.03);  // 设置字体大小
        legend->Draw();
    }
    c->Update();
}

void draw_tpc(const char* fileName = "Unpack.root") {
    gStyle->SetLabelFont(62, "XYZ");
    gStyle->SetTitleFont(62, "XYZ");
    gStyle->SetTitleFont(62, "");
    gStyle->SetStatFont(62);
    gStyle->SetLabelSize(0.05, "XYZ");
    gStyle->SetTitleSize(0.05, "XYZ");
    // gStyle->SetPalette(kRainBow); 
    gStyle->SetPalette(1); 
    //参数2和3代表长和宽
    new DataViewer(gClient->GetRoot(), 1600, 800, fileName);  
}

